package cn.xxi.ai.knowledge.constant;

import cn.xxi.ai.config.MinioProperties;
import cn.xxi.ai.util.SpringContextHolder;
import dev.langchain4j.community.model.dashscope.QwenEmbeddingModel;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.googleai.GoogleAiEmbeddingModel;
import dev.langchain4j.model.ollama.OllamaEmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.milvus.MilvusEmbeddingStore;
import io.milvus.client.MilvusServiceClient;
import lombok.Getter;

import java.util.Collections;
import java.util.List;
import java.util.Objects;

/**
 * 嵌入模型枚举，定义不同的Embedding模型及其提供方
 *
 * @author yangzhen
 */
@Getter
public enum EmbeddingModelEnum {

    // Ollama平台的中文大模型
    QUENTINZ_BGE_LARGE_ZH_V1_5_LATEST("quentinz/bge-large-zh-v1.5:latest", "ollama"),

    // DashScope平台的文本嵌入模型
    TEXT_EMBEDDING_V3("text-embedding-v3", "dashscope"),
    TEXT_EMBEDDING_V2("text-embedding-v2", "dashscope"),
    TEXT_EMBEDDING_V1("text-embedding-v1", "dashscope"),

    // Gemini平台的嵌入模型
    EMBEDDING_001("embedding-001", "gemini"),
    TEXT_EMBEDDING_004("text-embedding-004", "gemini");

    private final String modelName;
    private final String provider;

    private EmbeddingModel model;
    private EmbeddingStore<TextSegment> store;

    EmbeddingModelEnum(String modelName, String provider) {
        this.modelName = modelName;
        this.provider = provider;
        initLazyComponents();
    }

    // 初始化模型和存储（模拟lazy）
    private void initLazyComponents() {
        MinioProperties properties = SpringContextHolder.getBean(MinioProperties.class);
        MilvusServiceClient milvusClient = SpringContextHolder.getBean(MilvusServiceClient.class);

        // 构建模型（根据provider）
        switch (provider) {
            case "ollama":
            case "openai":
                this.model = OllamaEmbeddingModel.builder()
                        .baseUrl(properties.getEndpoint())
                        .modelName(modelName)
                        .build();
                break;
            case "dashscope":
                this.model = QwenEmbeddingModel.builder()
                        .apiKey(System.getenv("DASHSCOPE_API_KEY"))
                        .modelName(modelName)
                        .build();
                break;
            case "gemini":
                this.model = GoogleAiEmbeddingModel.builder()
                        .apiKey(System.getenv("GEMINI_API_KEY"))
                        .modelName(modelName)
                        .build();
                break;
            default:
                throw new IllegalArgumentException("Unsupported provider: " + provider);
        }

        // 构建 Milvus 存储
        this.store = MilvusEmbeddingStore.builder()
                .milvusClient(milvusClient)
                .collectionName(provider + "_" + name().toLowerCase())
                .dimension(Objects.requireNonNull(model).dimension())
                .build();
    }

    /**
     * 嵌入并存储多个文本段
     */
    public void embedAndStore(List<TextSegment> textSegments, List<String> segmentIds) {
        var response = model.embedAll(textSegments);
        store.addAll(segmentIds, response.content(), textSegments);
    }

    /**
     * 嵌入并存储单个文本段
     */
    public void embedAndStore(TextSegment textSegment, String segmentId) {
        embedAndStore(Collections.singletonList(textSegment), Collections.singletonList(segmentId));
    }
}
